#ifndef STRUCTMESHPART_HPP
#define STRUCTMESHPART_HPP

#include "typeAliases.hpp"
#include "mpiHelpers.hpp"
#include "meshTypes.hpp"
#include "crtpBase.hpp"
#include "structMeshOverlap.hpp"
#include "field.hpp"

namespace hms {

/**
 * @brief For dynamic load balancing, three types of loads can be provided:
 * - one uniform value per node (e.g. the CPU load as reported by the OS)
 * - one value for each domain cell on the node holding these cells
 * - one value for each domain cell of the global mesh available on the master node
 * 
 */
enum class PartLoad {
	perNode,
	perCell,
	perCellRecomp,
};

template<typename PartType>
class MeshPartBase : public CrtpBase<PartType>
{};

/**
 * @brief Represents a mesh partition for a structured mesh type (i.e.
 * UniMesh, RectMesh or StructMesh). Offers functionality for (re-)distributing
 * meshes and fields, retrieving boundary values from neighbours and
 * rebalancing.
 * 
 * @tparam MeshType 
 */
template<typename MeshType>
class StructMeshPart : public MeshPartBase<StructMeshPart<MeshType>>
{
	static_assert( std::is_base_of_v<StructMeshBase<MeshType>, MeshType>,
		"Mesh must be structured!" );

public:
	using SizeType = typename MeshTraits<MeshType>::SizeType;

private:
	MPI_Comm_H m_comm;
	
	Array2i m_meshSizeGlobal {0,0};
	Array2i m_nParts {1,1};
	Array2i m_offset {0,0};

	MeshType m_mesh;

	std::array<std::vector<OverlapSegment>,4> m_neighbours;

	Array2Xi m_oldSizes;

	Array2Xi m_boundaryPatchLimits;

public:
	StructMeshPart( MPI_Comm_H comm = MPI_COMM_NULL );

	StructMeshPart(
		StructMeshBase<MeshType>&& mesh,
		MPI_Comm_H comm = MPI_COMM_WORLD
	);

	StructMeshPart(
		StructMeshBase<MeshType>&& mesh,
		Array2i offset,
		MPI_Comm_H comm = MPI_COMM_WORLD
	);

	
	auto meshSizeGlobal() const -> const Array2i&;
	auto nParts        () const -> const Array2i&;
	auto offset        () const -> const Array2i&;
	auto oldSizes      () const -> const Array2Xi&;
	auto meshSize      () const -> SizeType;
	auto band          () const -> Index;
	auto indexInBand   () const -> Index;
	auto comm          () const -> const MPI_Comm_H&;
	auto comm( MPI_Comm_H comm )-> void;
	
	auto mesh() const -> const MeshType&;
	auto mesh() -> MeshType&;
	

	auto neighbours( BoundaryLocation ) const
		-> const std::vector<OverlapSegment>&;

	bool atBoundary( BoundaryLocation ) const;

	auto getOldOffsets() const -> Array2Xi;

	void assignAndDistribute( MeshType&& );
	void assignLocal( MeshType&& mesh, Array2i offset );

	void rebalanceMeshSizes(PartLoad loadType, const Ref<const ArrayXs>& load);
	void printLoad         (PartLoad loadType, const Ref<const ArrayXs>& load) const;
	void redistributeMesh( const Ref<const Array2Xi>& meshSizes );

	void distributeField( Field<MeshType>& ) const;
	void redistributeField( Field<MeshType>&, bool recomposeFirst = false ) const;
	void redistributeFields(
		std::vector<Field<MeshType>>& fields,
		const Ref<const Array2Xi>& oldOffsets,
		const std::vector<OverlapBlock>& overlappingBlocks
	) const;

	void getNeighbourValues( Field<MeshType>& ) const;

	void boundaryPatchLimits( Array2Xi&& );
	auto boundaryPatchLimits() const -> const Array2Xi&;
	void clearBoundaryPatchLimits();

private:
	auto neighbours( BoundaryLocation ) -> std::vector<OverlapSegment>&;
	void findNeighbours(
		const Ref<const Array2Xi>& allSizes,
		const Ref<const Array2Xi>& allOffsets
	);

	Array2Xi gatherMeshSizes() const;
	void redistributeField_distr( Field<MeshType>& field ) const;
	void redistributeField_centr( Field<MeshType>& field ) const;
};

}

#endif
